import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import io
import base64

class DataHandler:
    def __init__(self, data):
        self.data = data
    
    def get_basic_info(self):
        """获取数据基本信息"""
        info = {
            'shape': self.data.shape,
            'columns': self.data.columns.tolist(),
            'dtypes': self.data.dtypes.to_dict(),
            'missing_values': self.data.isnull().sum().to_dict(),
            'numeric_columns': self.data.select_dtypes(include=[np.number]).columns.tolist(),
            'categorical_columns': self.data.select_dtypes(include=['object']).columns.tolist(),
            'description': self.data.describe().to_dict()
        }
        return info
    
    def create_correlation_heatmap(self):
        """创建相关性热图"""
        numeric_data = self.data.select_dtypes(include=[np.number])
        
        if numeric_data.shape[1] < 2:
            return None
        
        plt.figure(figsize=(10, 8))
        sns.heatmap(numeric_data.corr(), annot=True, cmap='coolwarm', center=0)
        plt.title('特征相关性热图')
        
        # 转换为base64字符串
        img_buffer = io.BytesIO()
        plt.savefig(img_buffer, format='png', bbox_inches='tight')
        img_buffer.seek(0)
        img_string = base64.b64encode(img_buffer.read()).decode()
        plt.close()
        
        return img_string
    
    def prepare_data(self, target_column):
        """准备训练数据"""
        # 简单的数据预处理
        data_clean = self.data.copy()
        
        # 删除缺失值
        data_clean = data_clean.dropna()
        
        # 分离特征和目标
        X = data_clean.drop(columns=[target_column])
        y = data_clean[target_column]
        
        # 只保留数值特征（简化处理）
        X_numeric = X.select_dtypes(include=[np.number])
        
        return X_numeric, y